#include <math.h>
#include <assert.h>
#include "../math/mathbase.h"
#include "bilateralfilter.h"

using namespace PBVP;

BilateralFilter::BilateralFilter()
{
	// Default values
	m_FilterRadius = 2;
	int mask_dim = m_FilterRadius*2 + 1;
	m_SpatMask = new double[mask_dim*mask_dim*mask_dim];

}

BilateralFilter::~BilateralFilter()
{
	delete[] m_SpatMask;
	
}

Reg3DVolumeData* BilateralFilter::filter(PBVP::Reg3DVolumeData *pData)
{
	unsigned int dim[3];
	pData->getDimension(dim);

	Reg3DVolumeData *pFiltered = new Reg3DVolumeData(dim);
	pFiltered->setNumOfVariables(1);
	float orig[3], span[3];
	pData->getOrigin(orig);
	pFiltered->setOrigin(orig);
	pData->getSpan(span);
	pFiltered->setSpan(span);

	float *fdata = new float[dim[0]*dim[1]*dim[2]];

	for (int i = 0; i < dim[0]*dim[1]*dim[2]; i++)
	{
		fdata[i] = 0;
	}

	double min = pData->getMinValue();
	double max = pData->getMaxValue();
	m_SigmaFunc = (max - min)/2;
	m_SigmaSpat = (m_FilterRadius -1) * std::max(std::max(span[0], span[1]), span[2]);
	printf("Function sigma = %f, Spatial sigma = %f\n", m_SigmaFunc, m_SigmaSpat);
	calcSpatialMask(pData);
	
	for (int k = 0; k < dim[2]; k++){
		for (int j = 0; j < dim[1]; j++) {
			for (int i = 0; i < dim[0]; i++)
			{
				int nv = i + j*dim[0] + k*dim[0]*dim[1];
				fdata[nv] = filterVertex(pData, dim, i, j, k);
			}
		}
	}

	pFiltered->setDataArray(fdata);
	printf("Filtered data: min = %f, max = %f\n", pFiltered->getMinValue(), pFiltered->getMaxValue());
	return pFiltered;
}

void BilateralFilter::filterIterative(PBVP::Reg3DVolumeData* &pData, unsigned int n)
{
	if (n <= 0) return;
	
	Reg3DVolumeData *preFilter = pData;
	Reg3DVolumeData *postFilter = filter(pData);
	n--;
	while (n > 0) {
		delete preFilter;
		preFilter = NULL;
		preFilter = postFilter;
		postFilter = filter(preFilter);
		n--;
	}
	delete preFilter;
	pData = postFilter;
}

float BilateralFilter::filterVertex(PBVP::Reg3DVolumeData *pData, unsigned int dim[], 
									int vx, int vy, int vz)
{
	float val = 0;
	double maskSum = 0;
	double x0 = pData->getValue(vx, vy, vz);

	int index = 0;
	for (int k = -m_FilterRadius; k <= m_FilterRadius; k++)
	{
		for (int j = -m_FilterRadius; j <= m_FilterRadius; j++)
		{
			for (int i = -m_FilterRadius; i <= m_FilterRadius; i++)
			{
				int v0 = Math::clamp(vx + i, 0, dim[0]-1);
				int v1 = Math::clamp(vy + j, 0, dim[1]-1);
				int v2 = Math::clamp(vz + k, 0, dim[2]-1);
				double x1 = pData->getValue(v0, v1, v2);
				double mask = exp(-0.5*(x0-x1)*(x0-x1)/(m_SigmaFunc*m_SigmaFunc))*m_SpatMask[index];
				// double mask = m_SpatMask[index];
				val += x1 * mask ;
				maskSum += mask;
				index++;
			}
		}
	}
	val = val / maskSum;

	return val;
}

void BilateralFilter::calcSpatialMask(PBVP::Reg3DVolumeData *pData)
{
	float span[3];
	pData->getSpan(span);
	
	for (int k = -m_FilterRadius; k <= m_FilterRadius; k++)
	{
		for (int j = -m_FilterRadius; j <= m_FilterRadius; j++)
		{
			for (int i = -m_FilterRadius; i <= m_FilterRadius; i++)
			{
				int m = maskId(i, j, k);
				double dst2 = (i*span[0])*(i*span[0]) + (j*span[1])*(j*span[1]) + (k*span[2])*(k*span[2]);
				m_SpatMask[m] = exp(-0.5*(dst2)/(m_SigmaSpat*m_SigmaSpat));
				// printf("mask[%d, %d, %d] = %f\n", i, j, k, m_SpatMask[m]);
			}
		}
	}
}

int BilateralFilter::maskId(int i, int j, int k)
{
	int sz = (m_FilterRadius << 1) + 1;
	int m = (i+m_FilterRadius) + (j+m_FilterRadius)*sz + (k+m_FilterRadius)*sz*sz;
	return m;
}